import numpy as np
from numpy.random import randint
import networkx as nx
import matplotlib.pyplot as plt

"""
一、基本模拟
1、案例描述
现在有3个筛子，D4：四面体，D6：六面体，D8：八面体
数字范围：
D4：1到4
D6：1到6
D8：1到8

假设骰子出现的概率均等均等概率
每个骰子的每一面出现的概率也均等
可以做以下简单的模拟：
操作1：选取一个筛子；
操作2：掷骰子拿到一个数字；
重复10次

2、结合隐马尔科夫模型分析
那么有两个序列：
（1）筛子的序列 -> 隐含状态链
（2）数字的序列 -> 可见状态链

这两个序列就构成了一个隐马尔科夫链，其的特点是：
（1）选取筛子的概率：HMM中说到的马尔可夫链其实是指隐含状态链，隐含状态（骰子）之间存在转换概率（transition probability）
（2）骰子本身每个面出现的概率：尽管可见状态之间没有转换概率，但是隐含状态和可见状态之间有一个概率叫做输出概率（emission probability）

以上，总结出隐马尔科夫模型六要素：
（1）隐含状态：集合
（2）隐含状态链：序列
（3）可见状态：集合
（4）可见状态链：序列
（5）转换概率：矩阵，另外转换或有时候称为转移在文本序列任务里，可以叫做“后接”
（6）输出概率：矩阵

其中核心的是1、3、5、6

可以看到，放在这个例子里，隐马尔科夫链的直观含义是很清晰的，但是用数学去描述的时候上升到抽象的高度，并不容易理解
"""

# dices = ["D4", "D6", "D8"]
#
# values_of_dices = {
#     "D4": np.arange(1, 5),
#     "D6": np.arange(1, 7),
#     "D8": np.arange(1, 9)
# }
#
# probs_of_values = {
#     "D4": [1/4 for i in range(4)],
#     "D6": [1/6 for i in range(6)],
#     "D8": [1/8 for i in range(8)]
# }
#
# epoch = 10
# show_indices = []
# show_values = []
# for i in range(epoch):
#     idx = randint(3)
#     dice = dices[idx]
#     value = np.random.choice(values_of_dices[dice])
#     show_indices.append(str(i) + "-" + str(dice))
#     show_values.append(str(i) + "-" +  str(value))
#
# # 绘图
# G = nx.Graph()
# for node in show_indices + show_values:
#     G.add_node(node)
#
# edges = []
# for i in range(epoch):
#     if i < epoch - 1:
#         edges.append((show_indices[i], show_indices[i + 1]))
#     edges.append((show_indices[i], show_values[i]))
#
# print("edges:", edges)
# r = G.add_edges_from(edges)
# nx.draw(G, with_labels=True, node_color="y")
# plt.show()

"""
二、HMM模型的三个应用方向
1、在刚才的例子里，我们对这个随机过程有着完全的信息：
（1）知道有几种骰子：隐含状态
（2）知道选取筛子的概率：转换概率
（3）知道骰子每个面出现的概率：输出概率
（4）能看到掷骰子的结果序列：可见状态。在这个状态下，很容易去写程序模拟出多次掷出筛子，获得一个数字序列（可见状态）。
但大多数时候，这四个信息里可能会缺失一些，这就要用算法去估算这些缺失信息。

2、HMM算法的三个类别
（1）评估问题：
隐含状态：已知
隐含状态链：已知
可见状态：已知
可见状态链：已知
转换概率：已知
输出概率：已知

直接计算当前结果出现的概率，以评估实际结果是否符合理论

（2）解码问题
隐含状态：已知
隐含状态链：求解
可见状态：已知
可见状态链：已知
转换概率：已知
输出概率：已知

（3）学习问题
隐含状态：已知
隐含状态链
可见状态
可见状态链：已知
转换概率：求解
输出概率
"""

# 解码，求解骰子序列
show_values = [1, 6, 3, 5, 2]
epoch = len(show_values)

dices = ["D4", "D6", "D8"]
num_dices = len(dices)

values_of_dices = {
    "D4": np.arange(1, 5),
    "D6": np.arange(1, 7),
    "D8": np.arange(1, 9)
}

probs_of_values = {
    "D4": [1/4 for i in range(4)],
    "D6": [1/6 for i in range(6)],
    "D8": [1/8 for i in range(8)]
}

# 前序遍历
show_value_probs = np.zeros((num_dices, epoch))
for i, value in enumerate(show_values):
    for j, dice in enumerate(dices):
        if value in values_of_dices[dice]:
            show_value_probs[j][i] = probs_of_values[dice][value - 1]
        else:
            show_value_probs[j][i] = 0


# def get_prob(show_value_probs, prob, row_idx, col_idx):
#     if col_idx > epoch - 1:
#         return prob
#     prob = prob * show_value_probs[row_idx][col_idx] * 1/3
#     col_idx += 1
#     # row_idx += 1
#     print("prob:", prob)
#     get_prob(show_value_probs, prob, row_idx, col_idx)

prob_list = []


def get_prob(show_value_probs, prob, row_idx, col_idx):
    if col_idx > epoch - 1:
        return prob
    prob = prob * show_value_probs[row_idx][col_idx] * 1/3
    col_idx += 1
    # row_idx += 1
    print("prob:", prob)
    get_prob(show_value_probs, prob, row_idx, col_idx)


# 暴力枚举：
guesses = []
print(show_value_probs)
# for i in range(epoch):
row_idx = 1
col_idx = 0
prob = 1
# get_prob(show_value_probs, prob, row_idx, col_idx)

# dice_guesses.extend(dices)
for i in range(1, epoch):
    current = []
    guesses = []
    if i == 1:
        prior = 1/3 * show_value_probs[:, i - 1]
        guess_prior = dices.copy()
    current_tmp = 1/3 * show_value_probs[:, i]
    for idx_prior, value in enumerate(prior):
        current.extend(list(value * current_tmp))
        for idx_dice in range(num_dices):
            guesses.append(guess_prior[idx_prior] + "->" + dices[idx_dice])
    prior = current.copy()
    guess_prior = guesses.copy()

print(len(current))
print(np.argmax(current))
print("np.max(current):", np.max(current))
print(guess_prior[30])
